import cv2
import os
import numpy as np
from PIL import Image, ImageDraw, ImageFont
#import RPi.GPIO as GPIO
import time

class FaceRecognizer:
    def __init__(self):
        # 初始化参数
        self.data_dir = 'data'
        self.names = []
        self.roles = []
        
        # 尝试读取name.txt，如果不存在则读取label_mapping.txt
        name_file = "name.txt"
        label_file = "label_mapping.txt"
        
        if os.path.exists(name_file):
            with open(name_file, 'r', encoding='utf-8') as f:
                lines = [line.strip() for line in f.readlines()]
                self.names = lines[:5]  # 前5行为用户名字
                self.roles = lines[5:]  # 第6行及以后为职责信息
        elif os.path.exists(label_file):
            with open(label_file, 'r', encoding='utf-8') as f:
                for line in f:
                    label, name_role = line.strip().split(':', 1)
                    # 将姓名和职业分开
                    if '_' in name_role:
                        name, role = name_role.split('_', 1)
                    else:
                        name, role = name_role, ""
                    self.names.append(name)
                    self.roles.append(role)
        else:
            print("错误：找不到用户信息文件（name.txt 或 label_mapping.txt）")
            return
        # 初始化识别器
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        self.recognizer.read("wenjian/trainer.yml")  # 加载训练好的模型

        # 加载人脸检测模型
        self.face_net = cv2.dnn.readNetFromCaffe(
            "wenjian/deploy.prototxt",
            "wenjian/res10_300x300_ssd_iter_140000.caffemodel"
        )

        # 加载字体
        self.font = ImageFont.truetype('wenjian/simhei.ttf', 20)

        # 初始化摄像头
        self.cap = cv2.VideoCapture(0)
        cv2.namedWindow('Face Recognition', cv2.WINDOW_NORMAL)
        cv2.resizeWindow('Face Recognition', 800, 480)
        cv2.moveWindow('Face Recognition', 0, 0)

    def detect_faces_dnn(self, frame):
        """检测人脸并返回坐标"""
        (h, w) = frame.shape[:2]
        blob = cv2.dnn.blobFromImage(cv2.resize(frame, (300, 300)), 1.0,
                                    (300, 300), (104.0, 177.0, 123.0))
        self.face_net.setInput(blob)
        detections = self.face_net.forward()
        faces = []
        for i in range(0, detections.shape[2]):
            confidence = detections[0, 0, i, 2]
            if confidence > 0.5:
                box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
                (startX, startY, endX, endY) = box.astype("int")
                faces.append((startX, startY, endX - startX, endY - startY))
        return faces

    def recognize_user(self):
        """识别用户并返回姓名和职业"""
        try:
            while True:
                ret, frame = self.cap.read()
                if not ret:
                    break
                
                frame = cv2.rotate(frame, cv2.ROTATE_180)
                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                frame_pil = Image.fromarray(frame_rgb)
                draw = ImageDraw.Draw(frame_pil)

                faces = self.detect_faces_dnn(frame)
                if len(faces) == 1:
                    x, y, w, h = faces[0]
                    draw.rectangle([x, y, x + w, y + h], outline=(0, 255, 0), width=2)
                    
                    # 识别当前人脸
                    face_roi = frame[y:y + h, x:x + w]
                    gray = cv2.cvtColor(face_roi, cv2.COLOR_BGR2GRAY)
                    id_, confidence = self.recognizer.predict(gray)
                    
                    # 显示识别结果
                    if confidence < 70:  # 置信度阈值
                        name = self.names[id_]
                        role = self.roles[id_] if id_ < len(self.roles) else ""
                        display_text = f"{name} {role}"
                        draw.text((x + w + 10, y), display_text, font=self.font, fill=(0, 255, 0))
                        
                        # 按下空格键返回结果
                        key = cv2.waitKey(30)
                        if key == 32:  # 空格键
                            return name, role
                    else:
                        draw.text((x + w + 10, y), "未识别到用户", font=self.font, fill=(255, 0, 0))
                
                cv2.imshow('Face Recognition', cv2.cvtColor(np.array(frame_pil), cv2.COLOR_RGB2BGR))
                
                key = cv2.waitKey(30)
                if key == 27:  # ESC退出
                    break

        except Exception as e:
            print(f"发生错误: {str(e)}")
        #finally:
           # self.cleanup()
        #return None, None

    '''def cleanup(self):
        """释放资源"""
        self.cap.release()
        cv2.destroyAllWindows()
        GPIO.cleanup()'''

if __name__ == "__main__":
    recognizer = FaceRecognizer()
    name, role = recognizer.recognize_user()
    if name:
        print(f"识别结果: 姓名={name}, 职业={role}")
    else:
        print("未识别到用户")